server : fix segfault on long system prompt #8987
Conversation
compilade
added
bugfix
|
Hello @compilade, I was trying to fix it as well, looks like you beat me to it 😅. Won't my fix be a little faster since it avoids the nested loop? Aside from that I have a small question. If clearing the batch doesn't clear the kv cache then this pr looks good to me👍 |
@pritam-dey3 Nice to see you were also trying to fix the problem you found :)
The copies of the system prompt tokens into the batch are done anyway, so I don't see why the nested loop would be slower. It's simply chunking the copy (Also, calling |
|
|
||
| for (int32_t i = 0; i < batch.n_tokens; i += n_batch) { |
There was a problem hiding this comment.
I guess there is some outdated logic in update_slots as well, because there is a similar loop there.
There was a problem hiding this comment.
The only problem I see in update_slots is when the batch size is smaller than the number of slots, in which case tokens are added to the batch anyway, which could be out-of-bounds.
llama.cpp/examples/server/server.cpp
Line 2002 in 06943a6
When the batch size is bigger than the number of slots, it does not seem to have the same problem as the system prompt, because there is already never more than n_batch tokens added to a batch:
llama.cpp/examples/server/server.cpp
Lines 2243 to 2254 in 06943a6
So yes, this loop:
llama.cpp/examples/server/server.cpp
Line 2313 in 06943a6
might be a bit outdated, but I think it's still relevant for the error handling logic:
llama.cpp/examples/server/server.cpp
Lines 2375 to 2376 in 06943a6
From the above I think what could be fixed is parallel generation with small batch sizes:
diff --git a/examples/server/server.cpp b/examples/server/server.cpp
index a1bbafbc..b2bf5bb5 100644
--- a/examples/server/server.cpp
+++ b/examples/server/server.cpp
@@ -753,13 +753,13 @@ struct server_context {
default_generation_settings_for_props = get_formated_generation(slots.front());
default_generation_settings_for_props["seed"] = -1;
- // the update_slots() logic will always submit a maximum of n_batch tokens
+ // the update_slots() logic will always submit a maximum of n_batch or n_parallel tokens
// note that n_batch can be > n_ctx (e.g. for non-causal attention models such as BERT where the KV cache is not used)
{
const int32_t n_batch = llama_n_batch(ctx);
// only a single seq_id per token is needed
- batch = llama_batch_init(n_batch, 0, 1);
+ batch = llama_batch_init(std::max(n_batch, params.n_parallel), 0, 1);
}
metrics.init();
* server : fix segfault on long system prompt * server : fix parallel generation with very small batch sizes * server : fix typo in comment
* server : fix segfault on long system prompt * server : fix parallel generation with very small batch sizes * server : fix typo in comment
Should fix #8975 (reported by @pritam-dey3).
The problem was introduced in #6017 by limiting the size of the
llama_batchused in the server ton_batchinstead ofn_ctx, while not updating the logic managing the system prompt processing to not put all the system prompt in thellama_batch.This fix limits how many tokens of the system prompt are put into the
llama_batchof the server at a time so that it never exceedsn_batch, while still allowing larger system prompts by splitting them across batches.cc @slaren